import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
from sklearn import tree
from sklearn.metrics import accuracy_score
import pickle

Xcols = ['eirp',
          'ap_from_ap_0_max_ant_rssi', 'ap_from_ap_0_mean_ant_rssi', 'ap_from_ap_0_sum_ant_rssi',
        # 'ap_from_ap_1_max_ant_rssi', 'ap_from_ap_1_mean_ant_rssi', 'ap_from_ap_1_sum_ant_rssi',
          'sta_from_ap_0_max_ant_rssi', 'sta_from_ap_0_mean_ant_rssi', 'sta_from_ap_0_sum_ant_rssi',
          'sta_from_ap_1_max_ant_rssi', 'sta_from_ap_1_mean_ant_rssi', 'sta_from_ap_1_sum_ant_rssi',
          'sta_to_ap_0_max_ant_rssi', 'sta_to_ap_0_mean_ant_rssi', 'sta_to_ap_0_sum_ant_rssi',
          'sta_to_ap_1_max_ant_rssi', 'sta_to_ap_1_mean_ant_rssi', 'sta_to_ap_1_sum_ant_rssi']

# 读取数据
df = pd.read_csv('training_merge_ReduceRssi.csv')
missing_values = df[Xcols].isnull().any(axis=1)
df = df[~missing_values]
print('data len', len(df))

# 分割特征和目标列
X = df[Xcols]
Y = df['mcs'] == 11
print(Y)

X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.1, random_state=42)

clf = DecisionTreeClassifier(max_depth = 8)
clf.fit(X_train, y_train)

# 保存模型
with open('mcsClass.pkl', 'wb') as f:
    pickle.dump(clf, f)

# 测试集看准确率
y_pred = clf.predict(X_test)
print(y_pred)
accuracy = accuracy_score(y_test, y_pred)
print("acc:", accuracy)

# 画决策树图
plt.figure(figsize=(30,8))
tree.plot_tree(clf, fontsize=7, feature_names=Xcols, class_names=['not 11', '11'])
plt.savefig('tree_high_dpi', dpi=100)